{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4c5e8cad-8a04-4952-b873-7de52998b72f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/afs/crc.nd.edu/user/y/ypeng4/.conda/envs/python3.10/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "`fused_weight_gradient_mlp_cuda` module not found. gradient accumulation fusion with weight gradient computation disabled.\n",
      "/afs/crc.nd.edu/user/y/ypeng4/.conda/envs/python3.10/lib/python3.10/site-packages/mmcv/__init__.py:20: UserWarning: On January 1, 2023, MMCV will release v2.0.0, in which it will remove components related to the training process and add a data transformation module. In addition, it will rename the package names mmcv to mmcv-lite and mmcv-full to mmcv. See https://github.com/open-mmlab/mmcv/blob/master/docs/en/compatibility.md for more details.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch\n",
    "import torchvision\n",
    "from torch import nn\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "import torch.nn.functional as F\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from utils import *\n",
    "__all__ = ['UNext']\n",
    "\n",
    "import timm\n",
    "from timm.models.layers import DropPath, to_2tuple, trunc_normal_\n",
    "import types\n",
    "import math\n",
    "from abc import ABCMeta, abstractmethod\n",
    "from mmcv.cnn import ConvModule\n",
    "import pdb\n",
    "\n",
    "\n",
    "\n",
    "def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:\n",
    "    \"\"\"1x1 convolution\"\"\"\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False)\n",
    "\n",
    "\n",
    "def shift(dim):\n",
    "            x_shift = [ torch.roll(x_c, shift, dim) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]\n",
    "            x_cat = torch.cat(x_shift, 1)\n",
    "            x_cat = torch.narrow(x_cat, 2, self.pad, H)\n",
    "            x_cat = torch.narrow(x_cat, 3, self.pad, W)\n",
    "            return x_cat\n",
    "\n",
    "class shiftmlp(nn.Module):\n",
    "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., shift_size=5):\n",
    "        super().__init__()\n",
    "        out_features = out_features or in_features\n",
    "        hidden_features = hidden_features or in_features\n",
    "        self.dim = in_features\n",
    "        self.fc1 = nn.Linear(in_features, hidden_features)\n",
    "        self.dwconv = DWConv(hidden_features)\n",
    "        self.act = act_layer()\n",
    "        self.fc2 = nn.Linear(hidden_features, out_features)\n",
    "        self.drop = nn.Dropout(drop)\n",
    "\n",
    "        self.shift_size = shift_size\n",
    "        self.pad = shift_size // 2\n",
    "\n",
    "        \n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            trunc_normal_(m.weight, std=.02)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "        elif isinstance(m, nn.Conv2d):\n",
    "            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
    "            fan_out //= m.groups\n",
    "            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))\n",
    "            if m.bias is not None:\n",
    "                m.bias.data.zero_()\n",
    "    \n",
    "#     def shift(x, dim):\n",
    "#         x = F.pad(x, \"constant\", 0)\n",
    "#         x = torch.chunk(x, shift_size, 1)\n",
    "#         x = [ torch.roll(x_c, shift, dim) for x_s, shift in zip(x, range(-pad, pad+1))]\n",
    "#         x = torch.cat(x, 1)\n",
    "#         return x[:, :, pad:-pad, pad:-pad]\n",
    "\n",
    "    def forward(self, x, H, W):\n",
    "        # pdb.set_trace()\n",
    "        B, N, C = x.shape\n",
    "\n",
    "        xn = x.transpose(1, 2).view(B, C, H, W).contiguous()\n",
    "        xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , \"constant\", 0)\n",
    "        xs = torch.chunk(xn, self.shift_size, 1)\n",
    "        x_shift = [torch.roll(x_c, shift, 2) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]\n",
    "        x_cat = torch.cat(x_shift, 1)\n",
    "        x_cat = torch.narrow(x_cat, 2, self.pad, H)\n",
    "        x_s = torch.narrow(x_cat, 3, self.pad, W)\n",
    "\n",
    "\n",
    "        x_s = x_s.reshape(B,C,H*W).contiguous()\n",
    "        x_shift_r = x_s.transpose(1,2)\n",
    "\n",
    "\n",
    "        x = self.fc1(x_shift_r)\n",
    "\n",
    "        x = self.dwconv(x, H, W)\n",
    "        x = self.act(x) \n",
    "        x = self.drop(x)\n",
    "\n",
    "        xn = x.transpose(1, 2).view(B, C, H, W).contiguous()\n",
    "        xn = F.pad(xn, (self.pad, self.pad, self.pad, self.pad) , \"constant\", 0)\n",
    "        xs = torch.chunk(xn, self.shift_size, 1)\n",
    "        x_shift = [torch.roll(x_c, shift, 3) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]\n",
    "        x_cat = torch.cat(x_shift, 1)\n",
    "        x_cat = torch.narrow(x_cat, 2, self.pad, H)\n",
    "        x_s = torch.narrow(x_cat, 3, self.pad, W)\n",
    "        x_s = x_s.reshape(B,C,H*W).contiguous()\n",
    "        x_shift_c = x_s.transpose(1,2)\n",
    "\n",
    "        x = self.fc2(x_shift_c)\n",
    "        x = self.drop(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "\n",
    "class shiftedBlock(nn.Module):\n",
    "    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n",
    "                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):\n",
    "        super().__init__()\n",
    "\n",
    "\n",
    "        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n",
    "        self.norm2 = norm_layer(dim)\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = shiftmlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            trunc_normal_(m.weight, std=.02)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "        elif isinstance(m, nn.Conv2d):\n",
    "            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
    "            fan_out //= m.groups\n",
    "            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))\n",
    "            if m.bias is not None:\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x, H, W):\n",
    "\n",
    "        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))\n",
    "        return x\n",
    "\n",
    "\n",
    "class DWConv(nn.Module):\n",
    "    def __init__(self, dim=768):\n",
    "        super(DWConv, self).__init__()\n",
    "        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)\n",
    "\n",
    "    def forward(self, x, H, W):\n",
    "        B, N, C = x.shape\n",
    "        x = x.transpose(1, 2).view(B, C, H, W)\n",
    "        x = self.dwconv(x)\n",
    "        x = x.flatten(2).transpose(1, 2)\n",
    "\n",
    "        return x\n",
    "\n",
    "class OverlapPatchEmbed(nn.Module):\n",
    "    \"\"\" Image to Patch Embedding\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768):\n",
    "        super().__init__()\n",
    "        img_size = to_2tuple(img_size)\n",
    "        patch_size = to_2tuple(patch_size)\n",
    "\n",
    "        self.img_size = img_size\n",
    "        self.patch_size = patch_size\n",
    "        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]\n",
    "        self.num_patches = self.H * self.W\n",
    "        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,\n",
    "                              padding=(patch_size[0] // 2, patch_size[1] // 2))\n",
    "        self.norm = nn.LayerNorm(embed_dim)\n",
    "\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            trunc_normal_(m.weight, std=.02)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "        elif isinstance(m, nn.Conv2d):\n",
    "            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
    "            fan_out //= m.groups\n",
    "            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))\n",
    "            if m.bias is not None:\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.proj(x)\n",
    "        _, _, H, W = x.shape\n",
    "        x = x.flatten(2).transpose(1, 2)\n",
    "        x = self.norm(x)\n",
    "\n",
    "        return x, H, W\n",
    "\n",
    "\n",
    "class UNext(nn.Module):\n",
    "\n",
    "    ## Conv 3 + MLP 2 + shifted MLP\n",
    "    \n",
    "    def __init__(self,  num_classes, input_channels=3, deep_supervision=False,img_size=224, patch_size=16, in_chans=3,  embed_dims=[ 128, 160, 256],\n",
    "                 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,\n",
    "                 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,\n",
    "                 depths=[1, 1, 1], sr_ratios=[8, 4, 2, 1], **kwargs):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder1 = nn.Conv2d(3, 16, 3, stride=1, padding=1)  \n",
    "        self.encoder2 = nn.Conv2d(16, 32, 3, stride=1, padding=1)  \n",
    "        self.encoder3 = nn.Conv2d(32, 128, 3, stride=1, padding=1)\n",
    "\n",
    "        self.ebn1 = nn.BatchNorm2d(16)\n",
    "        self.ebn2 = nn.BatchNorm2d(32)\n",
    "        self.ebn3 = nn.BatchNorm2d(128)\n",
    "        \n",
    "        self.norm3 = norm_layer(embed_dims[1])\n",
    "        self.norm4 = norm_layer(embed_dims[2])\n",
    "\n",
    "        self.dnorm3 = norm_layer(160)\n",
    "        self.dnorm4 = norm_layer(128)\n",
    "\n",
    "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]\n",
    "\n",
    "        self.block1 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.block2 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[2], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.dblock1 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.dblock2 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0],\n",
    "                                              embed_dim=embed_dims[1])\n",
    "        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1],\n",
    "                                              embed_dim=embed_dims[2])\n",
    "\n",
    "        self.decoder1 = nn.Conv2d(256, 160, 3, stride=1,padding=1)  \n",
    "        self.decoder2 =   nn.Conv2d(160, 128, 3, stride=1, padding=1)  \n",
    "        self.decoder3 =   nn.Conv2d(128, 32, 3, stride=1, padding=1) \n",
    "        self.decoder4 =   nn.Conv2d(32, 16, 3, stride=1, padding=1)\n",
    "        self.decoder5 =   nn.Conv2d(16, 16, 3, stride=1, padding=1)\n",
    "\n",
    "        self.dbn1 = nn.BatchNorm2d(160)\n",
    "        self.dbn2 = nn.BatchNorm2d(128)\n",
    "        self.dbn3 = nn.BatchNorm2d(32)\n",
    "        self.dbn4 = nn.BatchNorm2d(16)\n",
    "        \n",
    "        self.final = nn.Conv2d(16, num_classes, kernel_size=1)\n",
    "\n",
    "        self.soft = nn.Softmax(dim =1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        B = x.shape[0]\n",
    "        ### Encoder\n",
    "        ### Conv Stage\n",
    "\n",
    "        ### Stage 1\n",
    "        out = F.relu(F.max_pool2d(self.ebn1(self.encoder1(x)),2,2))\n",
    "        t1 = out\n",
    "        ### Stage 2\n",
    "        out = F.relu(F.max_pool2d(self.ebn2(self.encoder2(out)),2,2))\n",
    "        t2 = out\n",
    "        ### Stage 3\n",
    "        out = F.relu(F.max_pool2d(self.ebn3(self.encoder3(out)),2,2))\n",
    "        t3 = out\n",
    "\n",
    "        ### Tokenized MLP Stage\n",
    "        ### Stage 4\n",
    "\n",
    "        out,H,W = self.patch_embed3(out)\n",
    "        for i, blk in enumerate(self.block1):\n",
    "            out = blk(out, H, W)\n",
    "        out = self.norm3(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "        t4 = out\n",
    "\n",
    "        ### Bottleneck\n",
    "\n",
    "        out ,H,W= self.patch_embed4(out)\n",
    "        for i, blk in enumerate(self.block2):\n",
    "            out = blk(out, H, W)\n",
    "        out = self.norm4(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "\n",
    "        ### Stage 4\n",
    "\n",
    "        out = F.relu(F.interpolate(self.dbn1(self.decoder1(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        \n",
    "        out = torch.add(out,t4)\n",
    "        _,_,H,W = out.shape\n",
    "        out = out.flatten(2).transpose(1,2)\n",
    "        for i, blk in enumerate(self.dblock1):\n",
    "            out = blk(out, H, W)\n",
    "\n",
    "        ### Stage 3\n",
    "        \n",
    "        out = self.dnorm3(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "        out = F.relu(F.interpolate(self.dbn2(self.decoder2(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        out = torch.add(out,t3)\n",
    "        _,_,H,W = out.shape\n",
    "        out = out.flatten(2).transpose(1,2)\n",
    "        \n",
    "        for i, blk in enumerate(self.dblock2):\n",
    "            out = blk(out, H, W)\n",
    "\n",
    "        out = self.dnorm4(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "\n",
    "        out = F.relu(F.interpolate(self.dbn3(self.decoder3(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        out = torch.add(out,t2)\n",
    "        out = F.relu(F.interpolate(self.dbn4(self.decoder4(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        out = torch.add(out,t1)\n",
    "        out = F.relu(F.interpolate(self.decoder5(out),scale_factor=(2,2),mode ='bilinear'))\n",
    "\n",
    "        return self.final(out)\n",
    "\n",
    "\n",
    "class UNext_S(nn.Module):\n",
    "\n",
    "    ## Conv 3 + MLP 2 + shifted MLP w less parameters\n",
    "    \n",
    "    def __init__(self,  num_classes, input_channels=3, deep_supervision=False,img_size=224, patch_size=16, in_chans=3,  embed_dims=[32, 64, 128, 512],\n",
    "                 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,\n",
    "                 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,\n",
    "                 depths=[1, 1, 1], sr_ratios=[8, 4, 2, 1], **kwargs):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.encoder1 = nn.Conv2d(3, 8, 3, stride=1, padding=1)  \n",
    "        self.encoder2 = nn.Conv2d(8, 16, 3, stride=1, padding=1)  \n",
    "        self.encoder3 = nn.Conv2d(16, 32, 3, stride=1, padding=1)\n",
    "\n",
    "        self.ebn1 = nn.BatchNorm2d(8)\n",
    "        self.ebn2 = nn.BatchNorm2d(16)\n",
    "        self.ebn3 = nn.BatchNorm2d(32)\n",
    "        \n",
    "        self.norm3 = norm_layer(embed_dims[1])\n",
    "        self.norm4 = norm_layer(embed_dims[2])\n",
    "\n",
    "        self.dnorm3 = norm_layer(64)\n",
    "        self.dnorm4 = norm_layer(32)\n",
    "\n",
    "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]\n",
    "\n",
    "        self.block1 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.block2 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[2], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.dblock1 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[1], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.dblock2 = nn.ModuleList([shiftedBlock(\n",
    "            dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=1, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
    "            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,\n",
    "            sr_ratio=sr_ratios[0])])\n",
    "\n",
    "        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0],\n",
    "                                              embed_dim=embed_dims[1])\n",
    "        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1],\n",
    "                                              embed_dim=embed_dims[2])\n",
    "\n",
    "        self.decoder1 = nn.Conv2d(128, 64, 3, stride=1,padding=1)  \n",
    "        self.decoder2 =   nn.Conv2d(64, 32, 3, stride=1, padding=1)  \n",
    "        self.decoder3 =   nn.Conv2d(32, 16, 3, stride=1, padding=1) \n",
    "        self.decoder4 =   nn.Conv2d(16, 8, 3, stride=1, padding=1)\n",
    "        self.decoder5 =   nn.Conv2d(8, 8, 3, stride=1, padding=1)\n",
    "\n",
    "        self.dbn1 = nn.BatchNorm2d(64)\n",
    "        self.dbn2 = nn.BatchNorm2d(32)\n",
    "        self.dbn3 = nn.BatchNorm2d(16)\n",
    "        self.dbn4 = nn.BatchNorm2d(8)\n",
    "        \n",
    "        self.final = nn.Conv2d(8, num_classes, kernel_size=1)\n",
    "\n",
    "        self.soft = nn.Softmax(dim =1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        B = x.shape[0]\n",
    "        ### Encoder\n",
    "        ### Conv Stage\n",
    "\n",
    "        ### Stage 1\n",
    "        out = F.relu(F.max_pool2d(self.ebn1(self.encoder1(x)),2,2))\n",
    "        t1 = out\n",
    "        ### Stage 2\n",
    "        out = F.relu(F.max_pool2d(self.ebn2(self.encoder2(out)),2,2))\n",
    "        t2 = out\n",
    "        ### Stage 3\n",
    "        out = F.relu(F.max_pool2d(self.ebn3(self.encoder3(out)),2,2))\n",
    "        t3 = out\n",
    "\n",
    "        ### Tokenized MLP Stage\n",
    "        ### Stage 4\n",
    "\n",
    "        out,H,W = self.patch_embed3(out)\n",
    "        for i, blk in enumerate(self.block1):\n",
    "            out = blk(out, H, W)\n",
    "        out = self.norm3(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "        t4 = out\n",
    "\n",
    "        ### Bottleneck\n",
    "\n",
    "        out ,H,W= self.patch_embed4(out)\n",
    "        for i, blk in enumerate(self.block2):\n",
    "            out = blk(out, H, W)\n",
    "        out = self.norm4(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "\n",
    "        ### Stage 4\n",
    "\n",
    "        out = F.relu(F.interpolate(self.dbn1(self.decoder1(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        \n",
    "        out = torch.add(out,t4)\n",
    "        _,_,H,W = out.shape\n",
    "        out = out.flatten(2).transpose(1,2)\n",
    "        for i, blk in enumerate(self.dblock1):\n",
    "            out = blk(out, H, W)\n",
    "\n",
    "        ### Stage 3\n",
    "        \n",
    "        out = self.dnorm3(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "        out = F.relu(F.interpolate(self.dbn2(self.decoder2(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        out = torch.add(out,t3)\n",
    "        _,_,H,W = out.shape\n",
    "        out = out.flatten(2).transpose(1,2)\n",
    "        \n",
    "        for i, blk in enumerate(self.dblock2):\n",
    "            out = blk(out, H, W)\n",
    "\n",
    "        out = self.dnorm4(out)\n",
    "        out = out.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()\n",
    "\n",
    "        out = F.relu(F.interpolate(self.dbn3(self.decoder3(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        out = torch.add(out,t2)\n",
    "        out = F.relu(F.interpolate(self.dbn4(self.decoder4(out)),scale_factor=(2,2),mode ='bilinear'))\n",
    "        out = torch.add(out,t1)\n",
    "        out = F.relu(F.interpolate(self.decoder5(out),scale_factor=(2,2),mode ='bilinear'))\n",
    "\n",
    "        return self.final(out)\n",
    "\n",
    "\n",
    "#EOF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "da6ac3b1-e366-4a32-8e3d-6124b548692b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 3, 256, 256])\n"
     ]
    }
   ],
   "source": [
    "model = UNext_S(num_classes=3)\n",
    "x = torch.rand((2,3,256,256))\n",
    "print(model(x).shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
